Multi-engine backends: pushdown DuckDB adapter, engine-neutral Arrow dataset, caching helpers#227
Draft
Mmoncadaisla wants to merge 62 commits into
Draft
Multi-engine backends: pushdown DuckDB adapter, engine-neutral Arrow dataset, caching helpers#227Mmoncadaisla wants to merge 62 commits into
Mmoncadaisla wants to merge 62 commits into
Conversation
Prototype of the "xarray-sql as the Xarray <-> engine translator" idea: the library owns two seams and nothing else. Seam 1 (register) attaches a lazy Dataset as a table on an engine's own connection; seam 2 (round-trip) turns any engine's Arrow result plus a template Dataset back into a labeled xr.Dataset. Dialects, geometry, H3, and optimizers stay with each engine and its extension ecosystem. - xarray_sql/backends/: adapter protocol + dispatch on connection type; DataFusion adapter delegates to the existing table provider, DuckDB adapter registers a re-scannable Arrow C-stream view (fresh lazy reader per scan: lazy AND re-queryable, no pushdown yet). - xarray_sql/roundtrip.py: engine-agnostic to_dataset() accepting DuckDB relations, pyarrow Tables/readers, or any __arrow_c_stream__ object; reuses the batches->Dataset core extracted from ds._materialize. - xql.register / xql.to_dataset exported at top level; [duckdb] extra. - docs/engines.md: the engine model, DuckDB usage, relation to duckdb-zarr (engine-native Zarr path; this adapter covers the rest of the xarray reader surface plus the labeled round-trip). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replace the default DuckDB registration object with XarrayPushdownDataset, a pyarrow.dataset.Dataset subclass (the pattern Lance uses for LanceDataset). DuckDB classifies it by isinstance and calls scanner(columns=..., filter=...) once per query, enabling: - projection pushdown: only the data variables a query mentions are loaded from storage; - chunk pruning: per-dimension shadow FileSystemDataset fragments carry each chunk's coordinate range as a partition_expression, so Arrow's guarantee simplification decides satisfiability for any predicate shape with no expression parsing on our side. One shadow per dimension keeps fragment counts at sum(n_d) instead of prod(n_d); - parallel production: surviving chunks are loaded by a bounded prefetch thread pool. DuckDB deletes pushed comparison conjuncts from its plan and never re-applies them, so the scanner always applies the exact expression via pyarrow Scanner; pruning is only an optimization. Filter-only columns absent from the projection are discovered by probing the expression against an empty table and widening on the miss. 10M-row benchmark vs the v1 stream: full scan 0.52s -> 0.035s, 1% time-filtered scan 0.24s -> 0.006s. A native-resolution bounding-box GROUP BY over a 9.13B-pixel cloud GeoTIFF answers in 0.8s on a plain duckdb connection (previously minutes). XarrayArrowStream stays as the dependency-light no-pushdown fallback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Bucketed two-level shadow pruning: axes with more chunks than the 1024-fragment fanout get a coarse bucket shadow refined lazily per surviving bucket, bounding pruning cost for finely partitioned datasets (e.g. hourly-chunked reanalysis time axes with hundreds of thousands of chunks). Refinement is skipped when a predicate keeps most buckets, where it cannot pay for itself. - Shadow datasets carry the full table schema so compound predicates referencing several columns bind and prune correctly. - Mixed-dimension datasets split into one table per dimension group (<name>_<dims>), sharing a single read of the dimension coordinates across sub-tables. - register() forwards adapter-specific kwargs (batch_size, prefetch on DuckDB; table_names on DataFusion). - Edge cases covered by tests: fully-pruned empty scans, LIMIT early termination, descending coordinate axes, uint8/string variables, 5000-chunk bucketed pruning loading exactly one chunk, kwarg forwarding. - duckdb extra pinned to >=1.4 (semantics verified against 1.5.4); engines doc gains production notes; pushdown benchmark added to benchmarks/. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Coordinate columns of a C-ordered partition are exact repeat/tile patterns of the dimension coords: dim k's flat column is its values each repeated prod(shape[k+1:]) times, tiled prod(shape[:k]) times. Building each column once per partition with those two sequential-write kernels and emitting batches as zero-copy Arrow slices replaces the per-batch division/modulo plus gather, making the pivot ~2.9x faster (53M -> 154M rows/s on a 10M-row, 3-dim dataset) with bitwise-identical output. Both engines benefit: iter_record_batches feeds the DataFusion table provider and the DuckDB pushdown scanner alike. The fast path holds the partition's full coordinate columns in memory (rows x 8 bytes x n_dims), so it is gated at 8M rows; larger partitions (e.g. single-time-step reanalysis) keep the O(batch_size) streaming path. The memory-profile test bands move accordingly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
rioxarray serializes GDAL reads behind a lock by default, capping any scan at single-stream speed regardless of the adapter's prefetch pool. lock=False measured 6x on full scans of a 9-billion-pixel cloud GeoTIFF (277s -> 43s) and makes remote reads as fast as a local copy. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Move XarrayPushdownDataset and XarrayArrowStream to
backends/pyarrow.py: the pushdown dataset is a real
pyarrow.dataset.Dataset, so it serves every consumer of that protocol,
not just DuckDB. The new public constructor xql.arrow_dataset(ds)
makes that explicit:
pl.scan_pyarrow_dataset(xql.arrow_dataset(ds)) # Polars, lazy
con.register("t", xql.arrow_dataset(ds)) # DuckDB
xql.arrow_dataset(ds).to_table(columns=..., filter=...) # pyarrow
Verified with Polars 1.42: scans are lazy, predicates and projections
push into the dataset (a filtered group-by read 1 of 20 chunks and 3
of 5 columns), results match xarray exactly, and Polars frames
round-trip through xql.to_dataset via the Arrow PyCapsule protocol.
The DuckDB adapter is now a thin registration shim over the shared
module. polars added to the test extra.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
get_fragments(filter=...) yields one fragment per source chunk, pruned by the same per-dimension shadow index the scanner uses. This is the protocol datafusion-python's register_dataset consumes (one DataFusion partition per fragment, filter pushdown marked Exact — safe because fragment scanners apply the expression row-exactly) and enables the Dask pattern from_map(lambda f: f.to_table().to_pandas(), ds.get_fragments()). Fragments carry a __dask_tokenize__ hook since the parent dataset is deliberately unpicklable. One arrow_dataset() object now serves DuckDB, Polars, DataFusion, Dask, and pyarrow itself. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Two fast paths in the batches->Dataset core (shared by every engine's round-trip): - Uniformly spaced axes (rasters, regular time steps, ascending or descending) locate each row with rint((value - origin) / step) instead of a per-row binary search plus argsort remap. 2.2x on a 25M-row chunk-ordered window (0.96s -> 0.43s); axes that are not affine within a quarter step keep the searchsorted path. - Results that form the complete grid in C order (ORDER BY'd scans, single-chunk windows) skip positional scatter entirely and reshape the value column. Both are detected, never assumed; sparse and arbitrarily ordered results fall back to the general scatter. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Registered xarray tables are virtual: every query re-streams the source. For statistics asked repeatedly, xql.materialize(con, name, query, order_by=...) pays the scan once into a native engine table (sorted so coordinate columns compress and zone maps prune), and xql.pyramid(con, name, table, aggs=..., base_cell=..., levels=...) builds a CARTO-tileset-style multi-resolution pre-aggregated cube: level 0 bins the source in a single pass, coarser levels roll up from the level below, so any zoom/extent query is a range scan over a small table. Aggregates are restricted to decomposable kinds (sum/count/min/max) so roll-ups stay exact; averages derive from sum + count at query time. Both helpers dispatch through a new EngineAdapter.run_sql seam and are tested identically on DuckDB connections and DataFusion contexts (the SQL they emit is restricted to the shared dialect subset; DataFusion INSERT is lazy, so run_sql collects). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Measured guidance in one place: parallel source reads (LIBERTIFF / lock=False, 11x on full raster scans; zarr async.concurrency), chunk sizing for the scan, adapter knobs (prefetch, batch_size), what pushdown does and does not cover, materialize/pyramid for repeated statistics, and the round-trip fast paths. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- NaN/NaT in a chunk's dimension coordinate poisoned the shadow guarantee into (dim >= NaN), which Arrow simplifies every predicate against as false — chunks with matching rows were silently pruned. Such spans now carry an always-true guarantee (unprunable). - _affine_axis accepted NaN-poisoned axes (NaN > tol is false), letting the affine scatter place values in wrong cells when a result's dim column contains NULLs. The acceptance test is now a <= .all() so NaN rejects and the positional searchsorted path handles it. - Projected scans omit dimension columns from the scan schema; the cftime coordinate-conversion loop assumed every dim had a schema field and raised KeyError on e.g. SELECT SUM(v) over a 360_day dataset. Dims absent from the schema are now skipped. - pyramid() rebinned float cell origins level-over-level, occasionally aliasing boundary points into a neighboring cell. Cells are now tracked as integer indices that halve exactly per level, with float origins kept as query labels. Also: docs notes from stress testing (Polars is_in float-literal pushdown caveat, per-thread DuckDB cursor re-registration pattern). Regression tests for all four fixes; battle-test matrix passed on duckdb 1.4.5 LTS and 1.5.4, 4-thread concurrency, dtype zoo, and a 600-query memory soak. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Polars (>= 1.40) passes batch_size through Dataset.to_batches to size its streaming-engine morsels; the scanner accepted the kwarg and dropped it, so consumers had no control over batch granularity. The emitted batch size matters beyond Polars: engines parallelize per record batch (DuckDB slices each batch into 2048-row vectors across threads, Acero schedules one filter task per batch), so it is the downstream-parallelism granule. Also pins two consumer-contract properties with tests: batch_size reaches the emitted batches through both scanner() and to_batches(), and the table schema never contains view types — one view-typed column disables DuckDB filter pushdown for the entire table (duckdb-python#227). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
count_rows() previously routed through a full scanner() scan, pivoting every surviving chunk to count rows the chunk grid already knows. Three-way split instead: pruned chunks contribute nothing, chunks whose coordinate ranges PROVE the filter true contribute their size arithmetically (no I/O), and only undecided boundary chunks are scanned — reading just the columns the filter references. Strictness is decided by Arrow itself, no expression parsing: a chunk with conjunctive guarantee G satisfies filter F everywhere iff G AND NOT F is unsatisfiable, which get_fragments(filter=~F) over per-chunk guarantee fragments answers (the Iceberg inclusive/strict evaluator pattern). Everything undecidable — NaN spans, data-variable predicates, >4096 survivors, unsupported expression shapes — falls back conservatively to an exact boundary scan. Also fixes columns=[] being treated as "all columns" (falsy-list bug): an explicitly empty projection is now a real projection, served by zero-column batches whose row counts are chunk arithmetic. Measured (10M-row synthetic, time chunked by 10): - unfiltered count: full scan -> 0.01 ms, zero chunks read - 5-day mid-chunk time range: exact count in 19 ms reading only the 2 boundary chunks (11 interior chunks proven arithmetically) - data-variable filter (no guarantee): still row-exact, all chunks scanned as before Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Finely chunked sources (an hourly-stepped time axis is one chunk per hour) pay one store round-trip per surviving chunk. With coalesce_rows=N, runs of consecutive surviving chunks along the most finely chunked dimension are merged into single isel reads of at most N rows, after pruning and per-run (no gaps are ever read). On Zarr sources the merged read fetches its member chunks through the store's own concurrent batching instead of one request per chunk through the prefetch pool. Scanner path only: get_fragments() keeps one fragment per source chunk so fragment consumers (DataFusion, dask) retain their parallelism granularity. Measured on ARCO-ERA5 over anonymous GCS (1.32M hourly time chunks, prefetch=16, coalesce_rows=8M), identical results both ways: - 1 day x Iberia bbox: 2.23s -> 1.18s (24 reads -> 4) - 1 week x full globe (174M rows): 9.82s -> 6.62s (168 reads -> 24) Peak RSS grows with prefetch x merged-block size as documented (0.7 GB -> 1.2 GB here); size coalesce_rows/prefetch together. Off by default: memory scales with the merged block, and local or coarsely chunked sources gain nothing. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Each scan previously created (and tore down) its own ThreadPoolExecutor. Beyond the per-scan setup cost, spawning OS threads from inside an engine's scan callback is exactly the embedded-engine hazard the ecosystem keeps rediscovering (pg_duckdb serializes all host calls; ParadeDB routes engine work through dedicated pools): thread startup contends with concurrent.futures' process-global shutdown lock and with the consumer's own pool management, observed as intermittent deadlocks when scans are driven from dask worker threads. The pool now lives on the dataset, its threads started at construction time — never inside an engine callback. Scans that stop early (LIMIT) cancel their queued loads instead of shutting the pool down. Single-block scans (a lazy round-trip window that maps onto one source chunk) skip the pool entirely. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
xql.to_dataset gains chunks= and coords=: data variables are reconstructed window by window on access, each window re-executing the engine's query narrowed to its coordinate range with the engine's own typed expression API (never rendered SQL text). Over a table registered through xarray-sql, the window's range predicate flows back into chunk pruning at the source, so accessing one output chunk reads only the source chunks it maps onto. The engine-specific surface of the existing DataFusion lazy path (expression building, per-dim distinct discovery, schema access) is extracted into LazyResultHandle implementations in the new xarray_sql.lazyscan module; SQLBackendArray and _build_lazy_scan are now engine-neutral, and the DataFusion wrapper path is byte-identical in behavior (39 pre-existing round-trip tests unchanged and green). Contiguous window requests become two-literal range predicates (engines can prune on them); stepped/fancy indexers fall back to explicit value lists, exact by construction. Engine support: - Polars LazyFrame/DataFrame: full chunked support; per-window fetches run on the streaming engine. Verified deadlock-free under threaded dask (6/6 stress runs) and correct on descending coordinates, stepped indexers, filtered and aggregated queries. - DataFusion DataFrame: unchanged, now also reachable through the engine-agnostic xql.to_dataset. - DuckDB relations: eager round-trip fully supported through a dedicated single engine thread (concurrent materialization of derived relations corrupts shared pending-query state — verified — so all handle calls are funnelled through one thread). Chunked reconstruction FAILS FAST with guidance instead of hanging: re-executing a relation that scans a Python-backed table while other threads start/stop deadlocks intermittently inside duckdb-python/CPython (~50% of runs, macOS/CPython 3.12; persists with SET threads=1, connection serialization, and pool pre-warming; Polars is clean under the identical topology). Reproducer kept in the working notes for an upstream report. coords="template" skips per-dim DISTINCT discovery when the result spans the template's full extent, making construction free of source reads for unfiltered scans on any engine. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Documents which engines support chunked reconstruction and why DuckDB relations fail fast (upstream duckdb-python deadlock, reproduced and isolated); states the memory bound (prefetch x pivoted-block-size) with the ARCO-ERA5 measurements that back it, and the coalesce_rows memory/latency tradeoff. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Asserts the SHAPE of the work — exactly which source chunks each query reads (via the scanner's iteration callback) and exact row counts — not just answers or timings, so a pruning/coalescing/fast-path regression that silently falls back to scanning everything fails loudly (the plan-shape-assertion pattern; benchmark-suite tripwires caught regressions plain tests missed in comparable projects). Measured this run (anonymous GCS, 1.32M-chunk hourly time axis): - day+bbox: 2.8s / exactly 24 chunk reads (4 reads coalesced, 1.6s) - week globe (174M rows): 16.6s / exactly 168 reads - count(*) over January (772M rows): 0.09s / ZERO reads (arithmetic) - Polars lazy round-trip: construction 0 reads; 1-day window compute reads exactly its 4 coalesced blocks Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
register(..., geometry=(x_dim, y_dim)) appends a geometry point column
derived from the coordinate dims, synthesized per batch at scan time —
the pivot's coordinate columns already carry the values, so the column
costs an annotation plus (for WKB) a vectorized 21-byte encode of the
rows actually scanned.
Two encodings, chosen per destination:
- "wkb" (default): geoarrow.wkb extension metadata + CRS. DuckDB >=1.2
with spatial loaded ingests the column as GEOMETRY('OGC:CRS84'), so
ST_Within(geometry, ...) works with no ST_Point construction in SQL.
- "point": GeoArrow-native separated coordinates; the struct children
ARE the coordinate arrays (no copy, no parse) for consumers that
execute on native layouts — verified with GeoPandas 1.x
GeoDataFrame.from_arrow, CRS carried through.
Documented sharp edge, measured: engines do not push ST_* functions
into the scan, so a geometry-only predicate defeats chunk pruning and
encodes every row (~29x slower than the paired form on a 10M-row
grid: 101ms bbox vs 2.9s ST_Within-only vs 118ms bbox+ST_Within).
The geospatial docs now state the idiom: bbox conjuncts for pruning,
geometry for exactness.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
bcba93a to
f9919b6
Compare
Polars translates float is_in literals imprecisely (reproduced on 1.42: is_in([<stored coordinate>]) matches zero rows). The handle now renders float value lists as OR-chains of degenerate is_between ranges, which compare exactly. Non-float dims keep is_in. A stepped lazy-round-trip window over non-representable float coordinates (linspace(-45, 45, 19)) previously scattered nothing and returned garbage; now row-exact. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The eager path materialized unconditionally: a billion-row result (or a sparse one whose dense coordinate-product grid dwarfs its Arrow payload) exhausted memory rather than erroring. max_result_bytes= now raises a clean ValueError with the running size at both danger points — while collecting the Arrow stream, and before allocating dense arrays (checked against the coordinate product, which is where sparse diagonals blow up). Error before OOM, never truncate; opt-in and unlimited by default. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…pruning The strictness analysis behind count_rows built one guarantee fragment per surviving chunk, capped at 4096 survivors; broader filters fell back to scanning every survivor (a January count on an hourly axis scanned 744 network chunks; step >= 100 over 1M chunks scanned ~1M). Replaced with hierarchical classification: each level buckets the surviving index lists into at most 4096 span-products, decides whole buckets at once with two guarantee-simplification passes (G AND NOT F unsat => proven; G AND F unsat => pruned), and recurses only into mixed cells. Per-chunk coordinate bounds are vectorized (np.minimum.reduceat, cached), so million-chunk axes classify in milliseconds. The prune side also refines the per-dimension pruning with cross-dimension information: paired ranges across dims ((A AND B) OR (C AND D)) no longer read the cross combinations. Measured: count over 1M single-row chunks with a near-universal filter 999,900 rows in 0.21s with ZERO chunk reads (previously scanned every survivor); cross-dim paired ranges read 2 boundary chunks instead of 4; 200 randomized differential checks against numpy ground truth all exact. NaN spans, non-numeric dims and simplifier-rejected expressions still land conservatively in the exact boundary scan. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The previous version used the single-lat-chunk fixture (no cross to refine) and hand-computed the wrong expected count; now differential against numpy on a grid chunked in both dims. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Count-based admission (prefetch=N blocks) under-uses the network with small chunks and overshoots memory with coalesced ones, because block sizes vary. prefetch_bytes= gates admission on estimated pivoted bytes in flight (rows x schema row width), the Lance io_buffer_size semantics; prefetch keeps bounding concurrency (thread count). With a byte budget, raising coalesce_rows no longer multiplies peak RSS. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Engines never push ST_* functions into the scan, so a geometry-only WHERE reads every chunk (documented, measured 29x). This helper renders the bbox range conjuncts from a geometry's envelope (a (xmin, ymin, xmax, ymax) tuple or anything with .bounds — shapely geometries qualify), making the bbox+geometry idiom one f-string instead of hand-copied bounds. Optional pad= margin for ST_DWithin-style predicates. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…mple) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ot streams The two results the re-execution path could not serve now work: spill=True streams the result exactly ONCE with bounded memory into a temporary Parquet file — through the engine handle where one exists (DuckDB spills on its dedicated engine thread, dodging the duckdb-under-worker-threads deadlock entirely; Polars uses its native streaming sink; DataFusion streams batch by batch) or straight from the Arrow stream for one-shot tables/readers — and the ordinary lazy reconstruction then runs over a Polars scan of that file, whose per-window predicates get Parquet row-group pruning. The temp file is deleted when the returned Dataset is garbage collected. Trade-off vs re-execution, by design: one full pass plus temporary disk instead of pay-per-window — the right shape when most of the result will be touched; Polars/DataFusion re-execution remains the default for window-at-a-time access. Verified: the previously deadlocking DuckDB chunked compute topology now passes 6/6 stress runs with repeated full computes; the source is read exactly once (chunk counter), and windows read only the spill file. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Engine-specific issues now sit under DuckDB / Polars / DataFusion content tabs (the Material tabbed pattern), so readers filter by their engine with one click; constraints that hold in every engine stay flat below, where a filter would mislead. Each entry keeps the symptom / scope / what-the-library-does shape. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The section lumped both helpers into one paragraph, leaving unclear what each builds. Now: a contrast table (any-query cache vs multi-resolution cube; CREATE TABLE AS SELECT vs raster overviews), one subsection each with a worked example, and pyramid gains the missing usage side — querying by level, range-filtering the float bin origins, joining on the exact integer indices. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
State what it actually is (CREATE OR REPLACE TABLE ... ORDER BY) and where the value lives: engine portability (DataFusion DDL is lazy and must be collected — the adapter handles it) and defaulting the sort-for-compression idiom. No overselling. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
materialize() wrapped four lines of CREATE OR REPLACE TABLE ... ORDER BY behind an API. That is engine SQL, and engine SQL belongs to the engine — the library owns register and round-trip, nothing between. The performance guide now shows the recipe directly (with the ORDER BY compression rationale and the one real quirk the wrapper hid: on DataFusion, DDL is a lazy plan that must be collected), and a recipe-regression test pins it on both engines. pyramid() stays: the multi-level rollup with exact integer cell indexing is genuine logic a recipe cannot carry. Its module is renamed materialize.py -> pyramid.py accordingly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A 4x4-pixel worked example whose entire 6-row output is shown inline — level 0 is the binned scan, level 1 is those cells added up — followed by the drill-down use case with measured numbers (365k-row cube from 16.7M pixels; overview reads 286 rows at ~1 ms, ~100x the raw GROUP BY locally, and the gap grows with source size). Seeing the rows explains the level/x_idx/x_bin columns better than describing them. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Same reasoning that removed materialize(), applied consistently: the aggregates dict-of-tuples was a Python DSL wrapping SQL fragments, at odds with the library's you-write-the-engine's-SQL thesis. The multi-resolution cube is two SQL statements and a loop; the performance guide now shows the recipe directly, keeps the 6-row worked example, and encodes the two hard-won details in prose — track cells as integer indices (float-origin rebinning aliases boundary points across parent cells) and range-filter the float bins. test_sql_recipes.py pins both documented recipes (caching and pyramid) verbatim on DuckDB and DataFusion, so the SQL cannot rot and the integer-index invariant stays tested. The library's public API is now exactly the two seams plus their supporting types. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…actionable Pyramid is gone at every level — helper (previous commit), documented recipe, recipe test, docstring mentions — leaving the caching recipe as the single stop-re-scanning pattern. The performance guide's round-trip section narrated internal fast paths (affine location, reshape vs scatter) without telling the reader anything to do; it is now two sentences of guidance: ORDER BY your dimension columns for large round-tripped results (grid-ordered arrives ~2x faster than the positional scatter unordered results pay), and reach for ds.sel when the question is not relational. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The README (and the docs homepage that snippets it), the examples page, and the geospatial intro still described a single-engine library. Updated in place, keeping their voice: What-is-this states the translate-data-not-queries position; the quickstart gains the one-call DuckDB registration and engine-agnostic round-trip next to the XarrayContext path; the how-it-works 2026 note explains that the same chunks-to-batches translation serves DuckDB and Polars through the pyarrow dataset protocol; the limitations paragraph points at the Known issues catalog instead of "TBD"; examples close with the same-tables-other-engines pointer. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A DuckDB user should not wade through Polars advice and vice versa. The guide now states up front that everything before the last section is engine-neutral (it tunes the shared scanner), and closes with a tabbed Per-engine notes section — the same filter-by-engine pattern as the Known issues page — collecting what is genuinely engine-specific: DuckDB connection-locality/cursors, chunk-ordered results, and the geometry/bbox pairing; Polars' single-threaded pull (tune prefetch, not Polars), batch sizing, and streaming collection; DataFusion's two registration paths (which knobs apply to which) and lazy DDL. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
One Coiled Function invocation per (workload, engine) cell on a shared
n2-standard-8 in us-central1 (in-region with the bucket), all SQL
engines reading through the same pushdown arrow_dataset (chunks
{'time': 1}, prefetch 16, coalesce_rows 8M); native xarray+dask
(chunks {'time': 24}) as the baseline. Wall-time rep policy (1 warm-up
+ median of 3 when the first rep is under 15s, else a single rep),
per-rep live logging, and a streamed results.jsonl of crash-safe
partials; errors are recorded as cells, never dropped.
Headline numbers (medians, 2026-07-14 run): the three SQL engines are
within ~3% of each other on scan-bound reductions -- week globe
(174M rows) 5.8/6.1/6.2s and month globe (772M rows) 24.9/25.6/24.9s
for DuckDB/Polars/DataFusion -- while native xarray stays ahead at
4.6s and 11.4s. count(*) over January: DuckDB 2.7s, Polars 4.3s
(11.5 GB peak RSS), DataFusion 16.6s. The 2-month regional monthly
climatology (GROUP BY lat, lon, month): DuckDB 54.2s, Polars 54.5s,
xarray 31.8s; DataFusion OOMed a 32 GB host on that cell (>26 GiB
unmanaged) and is recorded as an error.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Document the pyarrow-dataset engine matrix (DuckDB, Polars, DataFusion, native xarray+dask baseline) under Results in docs/geospatial.md: what each of the five workloads measures with its exact SQL and xarray expressions, the n2-standard-8 results table and condensed machine/ version spec, and a placeholder for the upcoming n2-standard-16 run. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
One table position with a switcher instead of two stacked tables; the 16-vCPU tab holds the pending placeholder until that run lands. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Superseded: the geospatial suite itself, run across engines and VM sizes, becomes the benchmark; its results will land in this page's Results section directly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
pa.array() returns a ChunkedArray instead of an Array for large numpy fixed-width unicode inputs (observed at ~3.3M rows of a string dimension coordinate tiled across a full partition), and RecordBatch.from_arrays rejects chunked input. Both iter_record_batches paths — the full-pivot fast path and the per-batch slow path — now flatten through a _as_single_array helper. Found by benchmarking: the forecast-skill geospatial case (05) streams a 2-model x 3.28M-row WeatherBench window whose "model" dim coord is a string; every pyarrow-protocol consumer (DuckDB, Polars, DataFusion) and the native reader crashed on it. Regression test included. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Make the nine-case geospatial suite engine-portable and measure it on three VM sizes, then publish the tables in the docs Results section. - benchmarks/geospatial/_engines.py: GEOBENCH_ENGINE facade. DataFusion keeps the original XarrayContext path when the native module exists (byte-for-byte unchanged) and falls back to SessionContext.register_dataset(xql.arrow_dataset(ds)) otherwise; DuckDB registers the same pushdown datasets; Polars runs the identical SQL via SQLContext with the query's window bounds also applied as scan-level expressions (its SQL TIMESTAMP literals compile to strptime casts that never reach the pyarrow scanner as filters). - Cases 01-06 route their register/query through the facade; SQL text, datasets, and correctness assertions are unchanged. 07/09 stay DataFusion-only (scalar UDFs); 08 stays on the original context (Earth-Engine-gated). - benchmarks/geospatial/engine_suite.py: Coiled Functions driver — one reused named VM per size run in parallel, run_perf.sh protocol (fresh process per rep, no warmup, 5 cold reps, answers asserted against the xarray reference), streamed results.jsonl, timestamped live logs, immediate cluster shutdown per VM. Replaces benchmarks/ engine_matrix.py (its coiled-function/jsonl patterns live on here). - docs/geospatial.md: per-VM tabs (e2-standard-8/16/32, us-central1) with the case x engine tables. Headline numbers (e2-8, medians of 5 cold reps): DuckDB leads every ARCO-ERA5 case through the shared pyarrow scan — 02 climatology 4.04s vs 4.36s Polars / 7.51s DataFusion-pyarrow (xarray 2.45s); 04 anomaly 6.09s vs 11.47s / 12.86s (xarray 5.10s); 06 zonal stats 4.71s vs 5.31s DataFusion (Polars: range-JOIN unsupported). Read-bound cases cluster tightly (01: 3.8-4.7s; 05: 1.7-2.3s). Extra vCPUs buy nothing - every case is single-stream read-bound, spelled out in the scope note along with the pyarrow-vs-native DataFusion caveat. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
VM size made no practical difference (all cases read-bound), so drop the e2-standard-16/-32 tabs and flatten the engine-suite results to the one e2-standard-8 table, with a condensed note that the larger sizes were measured and the full per-size tables live in xarray-sql-notes. Cases 07-09 were skipped on the Coiled VMs (no Earth Engine auth); their rows now come from the headline run as a clearly-marked adjacent table (original e2-standard-8 GCE VM with Earth Engine access, DataFusion native path). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Remove EngineAdapter.run_sql and both adapter implementations: its only callers, materialize() and pyramid(), were removed earlier on this branch (the docstring still pointed at xql.materialize). - roundtrip: drop the unused _apply_template import and _guarded's unused schema parameter; fold the five-way Arrow stream probe that _result_to_batches and _stream_to_parquet each spelled out into a shared _open_stream helper. - roundtrip docstrings: stop listing DuckDB relations as satisfying chunks= — they refuse the chunked path (DuckDBHandle.supports_chunked) and need spill=True. - lazyscan: share DuckDBHandle's version-probed Arrow reader construction (_to_arrow_reader), previously duplicated between fetch() and spill_parquet(). - pyarrow: factor _guarantee_shadow out of the fmt/fs FileSystemDataset construction repeated in _DimShadow._build and _strict_partition; factor _surviving and _outer_rows out of the surviving-chunk lookup and unresolved-dims row product repeated across the combos, coalescing, and strictness paths. - era5 benchmark: drop the stale caveat about the removed 4096-survivor strictness cap; _engines: collapse _pandas_to_dataset's redundant single-dim branch. - tests: hoist six in-function XarrayPushdownDataset imports to the top; share the four identical hourly-grid datasets (_hourly_grid) and the DuckDB spatial setup (spatial_con fixture); replace the vacuous inflight assertion in the prefetch_bytes test (its counter never decremented, so the bound held trivially) with exact read-count checks; drop an unused pandas import. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
A contiguous positional window was translated to a ('range', min, max)
spec whenever the window's own values were monotonic. That is only
exact when the entire coordinate array is monotonic: template coords
are used verbatim, and over an unsorted coordinate the range admits
values at positions outside the window, which the scatter then writes
over requested cells (reproduced: coord [102, 7, 55, 23], positions
1:3 -> range [7, 55] also matches 23 -> corrupted output).
The monotonicity of each dim's full coordinate array is now computed
once in SQLBackendArray.__init__ and passed into _dim_spec; windows
over non-monotonic coordinates fall back to explicit value lists.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The pre-spawn submitted N no-op tasks, but an idle thread absorbs several of them, so only a few OS threads actually started; the rest spawned later inside an engine's scan callback — the interleaving the shared pool exists to prevent. Each pre-spawn task now parks on a Barrier(N+1) that the constructor joins last, so no thread can take a second task and all N threads provably exist before __init__ returns. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Two eager-path routes escaped the budget: the Polars LazyFrame fallback collected the whole result inside the engine before the first batch reached the guard, and the to_arrow_table() branch never checked at all. Re-executable handles now expose stream(columns), yielding Arrow batches incrementally (DataFusion via execute_stream, Polars via collect_batches on the streaming engine), and the guarded collection consumes that stream when a budget is set — on Polars older than 1.33 (no collect_batches) it refuses up front instead of erroring after the memory is spent. The to_arrow_table() branch checks the materialized table's nbytes against the budget and is documented as materialize-then-check; results that also have a handle stream through it instead when a budget is set. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The streaming collects relied on the engine keyword without a version floor anywhere: pyproject's test extra accepted any polars, and a pre-1.25 install would raise TypeError mid-query. The test extra now pins polars >= 1.33 (engine="streaming" landed in 1.25; LazyFrame.collect_batches, used for streamed max_result_bytes enforcement, in 1.33), and PolarsHandle routes every collect through a _collect_streaming helper that falls back to the plain in-memory collect when the keyword is unknown. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
pa.binary() carries int32 offsets; past ~102M points per batch the final offset (n * 21) overflows and the array would build silently corrupt. The pivot's batch_size cap keeps this unreachable in normal use, so the encoder now raises a clear ValueError for direct callers instead of widening the storage to large_binary, which DuckDB's GEOMETRY ingestion does not consume. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The exact float translation built a left-deep OR chain of degenerate is_between ranges, which Polars plans quadratically in the number of values (measured: 0.7s at 1000 values, 2.7s at 2000, minutes at 10000). any_horizontal expresses the same disjunction as one flat node and stays exact for non-representable float coordinates. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The long-lived prefetch pool deliberately survives individual scans, but nothing ever stopped it, so every discarded dataset leaked its worker threads for the life of the process. A weakref.finalize on the dataset now shuts the pool down (wait=False, cancel_futures=True) once the dataset is collected; live scans keep the dataset alive through their generator closures, so in-flight work is unaffected. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Every DuckDBHandle starts a dedicated single-thread executor for engine calls and never stopped it, leaking one thread per discarded handle for the life of the process. A weakref.finalize on the handle now shuts the runner down (wait=False, cancel_futures=True) once the handle is collected; the callback is bound to the executor, so the finalizer holds no reference back to the handle. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
_check_dense_size counted per-dim distinct values by pouring every batch through to_pylist into a Python set — a per-row Python object allocation on results that can be millions of rows. pyarrow's unique kernel over a chunked view of the same columns computes the identical counts without leaving Arrow. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The limitations-page restructure renamed its first heading, orphaning engines.md's #upstream-issues anchor; and two geospatial.md links pointed at repo files outside docs/ that the site cannot serve — the README link now goes to the docs home (which embeds it), the benchmarks link to the repository path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The repo moved to xqlsystems: seven old-issue links 404'd outright and site_url pointed at a dead GitHub Pages host. Links that escape docs/ (benchmarks, perf_tests) and the README's repo-relative links 404 on the published site (index.md snippets the README) — now absolute. Also: move a Python comment out of a SQL string in the README example, fix a dangling colon and two typos. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DataFusion is the default engine; tabs now read DataFusion, DuckDB, Polars. Also add the missing blank line before two headings that Python-Markdown was rendering as body text. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
engines.md told users to pip install xarray-sql[duckdb] but the Polars section assumed polars was already present. Add polars>=1.33 as an extra, document the install, and quote both extras for zsh. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This branch prototypes xarray-sql as the Xarray ↔ engine translator: the library owns exactly two seams — register (a lazy
xarray.Datasetbecomes a table on an engine's own connection) and round-trip (any engine's Arrow result becomes a labeledxr.Datasetagain). SQL dialects, geometry functions, H3, and optimizers stay with each engine and its extension ecosystem. No transpiler, no unified dialect.flowchart TB A["Zarr · NetCDF · GRIB · GeoTIFF · Earth Engine"] --> B["lazy xarray.Dataset"] B --> C["seam 1 — register<br/>xql.register(con, name, ds) / xql.arrow_dataset(ds)"] C --> D["DuckDB<br/>+ spatial, h3, …"] C --> E["Polars"] C --> F["DataFusion"] C --> G["Dask"] D --> H["seam 2 — round-trip<br/>xql.to_dataset(result, template=ds)"] E --> H F --> H H --> I["labeled xr.Dataset — SQL in, array out"]One object serves every engine:
xql.arrow_dataset(ds)is a realpyarrow.dataset.Dataset(the pattern Lance uses forLanceDataset), so DuckDB registers it, Polars scans it viascan_pyarrow_dataset, DataFusion consumes it viaregister_dataset(through the fragments API), and Dask maps overget_fragments(). Ibis works throughibis.duckdb.from_connectionwith zero code.How the pushdown scan works
flowchart TB Q["engine calls scanner(columns, filter)"] --> P["prune chunks<br/>per-dimension shadow fragments;<br/>Arrow guarantee simplification decides satisfiability"] Q --> J["project<br/>only referenced variables are read"] P --> C["coalesce (opt-in)<br/>merge consecutive surviving chunks<br/>into single reads"] C --> L["prefetch pool<br/>bounded by prefetch (threads)<br/>and prefetch_bytes (memory)"] J --> L L --> X["exact filter<br/>pyarrow applies the pushed expression row-exactly"] X --> R["Arrow batches → engine"]Three properties worth calling out:
FileSystemDatasetfragments carry each chunk's coordinate range as apartition_expression(their paths are never opened). Sound for every predicate shape — equality,OR,IN,NOT.count(*)never scans: unfiltered counts are chunk arithmetic, and coordinate-range counts read at most the boundary chunks.The chunked round-trip is engine-generic
xql.to_dataset(result, chunks=...)reconstructs a query result as a chunked, lazyxr.Dataset— each window re-executes the engine's query narrowed to its coordinate range, which flows back into chunk pruning at the source.spill=Trueprovides the alternative one-pass shape: stream once (bounded memory) into a temporary Parquet file that windows re-execute against.flowchart TB R["xql.to_dataset(result, ...)"] --> K{"chunks=?"} K -- "None (default)" --> E["eager: materialize once<br/>(max_result_bytes= guards the stream<br/>and the dense grid)"] K -- "mapping / auto / inherit" --> SP{"spill=?"} SP -- "False (default)" --> RX["re-execution<br/>Polars & DataFusion results"] SP -- "True / directory" --> SPL["one-pass spill → temp Parquet<br/>the chunked path for DuckDB relations<br/>and one-shot Arrow streams"]coords="template"skips coordinate discovery entirely for full-extent results: on ARCO-ERA5 (1.32M hourly chunks) it builds a lazy view over a 1.37-trillion-row table in ~0.3 s with zero source reads; a one-day window then computes in ~2 s reading only the source chunks under it.Performance
Measured on a public 9-billion-pixel cloud-optimized GeoTIFF and a 10M-row synthetic benchmark (
benchmarks/duckdb_pushdown.py):GROUP BY, native res, full table registereddocs/performance.md)AVGscan (vs v1 stream)to_dataset)Peak scan memory is a contract, not an accident: bounded by
prefetch × pivoted-block-sizeregardless of data scanned (a 772M-row month-scale ARCO-ERA5 aggregation peaks at the same ~0.75 GB RSS as the week-scale scan). The pivot fast path (whole-partition repeat/tile coordinate columns, zero-copy batch slices) and the round-trip fast paths (affine-axis scatter, grid reshape) live in shared code, so the existing DataFusion engine benefits equally.Beyond microbenchmarks,
benchmarks/geospatial/runs nine staples of geospatial/climate analysis (NDVI, climatology, anomalyJOIN, forecast skill vs WeatherBench 2, raster×vector zonal stats, PROJ-UDF reprojection, weight-table regridding, warp) in SQL against real cloud datasets, each asserted against an xarray reference;engine_suite.pyrepeats the portable cases across DataFusion, DuckDB, and Polars on GCE VMs. The write-up isdocs/geospatial.md.What's in the branch
xarray_sql/backends/— adapter protocol with dispatch on connection type; DataFusion adapter delegates to the existing table provider; DuckDB adapter registers the pushdown dataset (mixed-dimension datasets split into one table per dim group, sharing coordinate reads).xarray_sql/backends/pyarrow.py—XarrayPushdownDataset(projection pushdown, shadow pruning, coalescing, prefetch, fragments API) andXarrayArrowStream(re-scannable C-stream fallback).xarray_sql/roundtrip.py+xarray_sql/lazyscan.py— engine-agnosticxql.to_datasetaccepting DuckDB relations, Polars frames, DataFusion DataFrames,pyarrowtables/readers, or any__arrow_c_stream__object; eager, re-executing, and spill-backed chunked reconstruction; metadata recovery from a template Dataset.xarray_sql/geometry.py—register(..., geometry=("x", "y"))derives a GeoArrow pointgeometrycolumn (WKB for DuckDB-nativeGEOMETRY, or GeoArrow-separated for GeoPandas/lonboard), CRS tagged;xql.bbox_conjunctsrenders prunable bbox predicates from any geometry's envelope.df.py/ds.py.docs/engines.md(the engine model, per-engine usage, support matrix),docs/performance.md(measured tuning guide),docs/limitations.md(known issues, pinned by tests),docs/geospatial.md(the relational-operations write-up with the benchmark results), reworkeddocs/examples.mdand README.[duckdb](duckdb>=1.4) and[polars](polars>=1.33) extras.Testing
275 tests, all green. Beyond unit coverage:
OR/IN, filter-column-outside-projection,LIMIT, empty results, NULL semantics) passes identically on duckdb 1.4.5 LTS and 1.5.4, plus the Polars battery on 1.42.max_result_bytesenforcement on eager collection and dense allocation.ChunkedArrayslipping into the pivot's fast path — all fixed with regression tests.Known limitations / follow-ups
chunks=raises immediately): re-execution from worker threads deadlocks inside duckdb-python 1.4–1.5 on CPython 3.12, sospill=Trueis the chunked path there. Documented with the full story indocs/limitations.md.XarrayPushdownDatasetsubclasses pyarrow's cythonDatasetwithout initializing the native base (as Lance does); dangerous inherited members are stubbed, and the contract is pinned by tests — re-verify on pyarrow upgrades.is_infloat literals imprecisely (reproducible without xarray-sql); documented, use range predicates. Polars also has no geometry types — thegeometrycolumn arrives as plain binary there.🤖 Generated with Claude Code